# pymym is a Python re-implementation of the serialization/deserialization
# functionality of mYm (mym.sourceforge.net) 

# Copyright (C) Suraj Gowda, 2012
# 
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301  USA
# 
# Author:                     Suraj Gowda 
# e-mail:                     surajgowda33@gmail.com 
# 
# This work is a derivative work of mYm (mym.sourceforge.net), and thus
# covered under the same license terms. 

import struct, pickle
from numpy import *
import numpy as np

ARRAY_HEADER = 'A'
CELL_HEADER = 'C'

def is_mym_data(s):
    if isinstance(s, basestring):
        return s[:3] == 'mYm'
    else:
        return False

def pack(obj, include_header=True):
    if include_header:
        header = 'mYm' + struct.pack('b', 0)
    else:
        header = ''
    if isinstance(obj, ndarray) or isinstance(obj, matrix):
        obj = array(obj)
        return header + pack_mym_array(obj)
    elif isinstance(obj, float) or isinstance(obj, int) or isinstance(obj, long):
        obj = array([[obj]])
        return header + pack_mym_array(obj)
    elif isinstance(obj, basestring):
        return header + pack_mym_str(obj)
    elif isinstance(obj, type(None)):
        return ''
    elif isinstance(obj, list):
        return header + pack_mym_cell(obj)
    elif isinstance(obj, dict): ## TODO
        raise NotImplementedError
    else:
        return pickle.dumps(obj)
        #raise NotImplementedError

def pack_mym_cell(s):
    header = CELL_HEADER
    ndim_str = struct.pack('i', 2)
    arr_shape_str = ''.join([struct.pack('i', x) for x in (1, len(s))])
    
    data = ''
    for item in s:
        packed_item = pack(item, include_header=False)
        data += struct.pack('i', len(packed_item))
        data += packed_item 
    return header + ndim_str + arr_shape_str + data


def pack_mym_str(obj):
    # make sure data is an array and not a matrix
    upsampled_str_ls = [0]*len(obj)*2
    upsampled_str_ls[::2] = obj

    header = ARRAY_HEADER
    ndim_str = struct.pack('i', 2) # Strings are fixed 1D objects
    arr_shape_str = ''.join([struct.pack('i', x) for x in (1,len(obj)) ])
    ctype_str = struct.pack('i', int(0)) # no complex-valued strings
    class_id_str = struct.pack('i', 4) # store as string 

    data = struct.pack('cb'*len(obj), *upsampled_str_ls) #''.join([struct.pack('d', x) for x in arr.reshape(-1, order='F')])
    full_header = header + ndim_str + arr_shape_str + class_id_str + ctype_str
    return full_header+ data
    #return header + ndim_str + arr_shape_str + class_id_str + ctype_str + data


def pack_mym_array(arr):
    """ Pack an array into binary string format used by mYm
    """
    # make sure data is an array and not a matrix
    arr = array(arr)

    header = ARRAY_HEADER
    ndim_str = struct.pack('i', ndim(arr))
    arr_shape_str = ''.join([struct.pack('i', x) for x in arr.shape])
    class_id_str = struct.pack('i', 6) # always store as double array

    complex_arr = np.any(np.iscomplex(arr))
    ctype_str = struct.pack('i', int(complex_arr))
    if complex_arr:
        real_data = ''.join([struct.pack('d', x) for x in np.real(arr).reshape(-1, order='F')])
        imag_data = ''.join([struct.pack('d', x) for x in np.imag(arr).reshape(-1, order='F')])
        data = real_data + imag_data
    else:
        data = ''.join([struct.pack('d', x) for x in arr.reshape(-1, order='F')])
    return header + ndim_str + arr_shape_str + class_id_str + ctype_str + data

def unpack(s, verbose=False):
    # check that s was inserted by mYm
    if s == None:
        return None
    elif is_mym_data(s):
        s = s[4:]
    elif isinstance(s, basestring) and len(s) == 0:
        return None
    else:
        try:
            return pickle.loads(s)
        except:
            return s

    # check that string is actually an array
    if s[0] == 'A':
        arr = unpack_mym_array(s[1:], verbose=verbose) 
        if shape(arr) == (1,1):
            return arr[0,0] 
        else:   
            return arr
    if s[0] == 'C':
        return unpack_mym_cell(s[1:]) 
    else:
        raise NotImplementedError("Type: %s" % s[0])
    
def unpack_mym_cell(s):
    idx = 0
    # get number of dimensions in array
    ndim = struct.unpack('i', s[idx:idx+4])[0]
    idx += 4
    arr_shape = struct.unpack('i'*ndim, s[idx:idx+4*ndim])

    idx += 4*ndim
    n_elem = prod(arr_shape)

    cell = [None]*n_elem
    for k in range(n_elem):
        nbytes = struct.unpack('i', s[idx:idx+4])[0]
        idx += 4

        cell[k] = unpack('mYm0'+s[idx:idx+nbytes])
        idx += nbytes

    return cell
    

def unpack_mym_array(s, verbose=False):
    """ Unpack an array stored as binary string 's' in MySQL db.
    Return as numpy array. 
    """

    idx = 0
    # get number of dimensions in array
    ndim = struct.unpack('i', s[idx:idx+4])[0]
    idx += 4
    arr_shape = struct.unpack('i'*ndim, s[idx:idx+4*ndim])

    if verbose:
        print "array shape"
        print arr_shape

    idx += 4*ndim

    # get mxClassID
    mxClassID = struct.unpack('i', s[idx:idx+4])[0]
    if mxClassID == 6:
        dtype = 'd'
        idx += 4
    elif mxClassID == 4:
        dtype = 'c'
        idx += 4
    else:
        print mxClassID
        return None

    # get complexity
    mxComplexity = struct.unpack('i', s[idx:idx+4])[0]
    idx += 4

    if dtype == 'c':
        arr_elems = array( struct.unpack( dtype*np.prod(arr_shape), s[idx::2] ) )
        return ''.join(arr_elems)
    elif dtype == 'd' and mxComplexity:
        all_elems = np.array( struct.unpack( dtype*np.prod(arr_shape)*2, s[idx:] ) )
        real_elems = all_elems[:len(all_elems)/2]
        imag_elems = all_elems[len(all_elems)/2:]
        return real_elems + 1j*imag_elems
        #return arr_elems#.reshape(arr_shape, order='F')
    else:
        arr_elems = np.array( struct.unpack( dtype*np.prod(arr_shape), s[idx:] ) )
        return arr_elems.reshape(arr_shape, order='F')
    
    
